package com.liyuxiang.graph.stream;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Sinks;

import java.util.Map;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class GraphStreamProcessor {

    private static final Logger logger = LoggerFactory.getLogger(GraphStreamProcessor.class);

    private final ExecutorService executor = Executors.newSingleThreadExecutor();
    private final ObjectMapper objectMapper;

    public GraphStreamProcessor(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    public void processStream(AsyncGenerator<NodeOutput> generator,
                              Sinks.Many<ServerSentEvent<String>> sink,
                              String threadId) {
        executor.submit(() -> {
            logger.info("开始处理流式输出 - 线程ID: {}", threadId);
            
            generator.forEachAsync(output -> {
                try {
                    String nodeName = output.node();
                    String content;
                    
                    logger.info("流式输出节点: {}, 类型: {}", nodeName, output.getClass().getSimpleName());
                    
                    if (output instanceof StreamingOutput streamingOutput) {
                        // 流式输出 - 逐步返回AI生成的内容
                        Map<String, Object> streamData = Map.of(
                            "type", "stream",
                            "node", nodeName,
                            "chunk", streamingOutput.chunk(),
                            "threadId", threadId
                        );
                        content = objectMapper.writeValueAsString(streamData);
                    } else {
                        // 普通节点输出 - 返回完整状态
                        Map<String, Object> nodeData = Map.of(
                            "type", "node",
                            "node", nodeName,
                            "data", output.state().data(),
                            "threadId", threadId
                        );
                        content = objectMapper.writeValueAsString(nodeData);
                    }
                    
                    ServerSentEvent<String> event = ServerSentEvent.<String>builder()
                        .id(String.valueOf(System.currentTimeMillis()))
                        .event("message")
                        .data(content)
                        .build();
                        
                    sink.tryEmitNext(event);
                    logger.debug("发送流式事件 - 节点: {}, 内容长度: {}", nodeName, content.length());
                    
                } catch (Exception e) {
                    logger.error("处理流式输出失败 - 线程ID: {}, 错误: {}", threadId, e.getMessage(), e);
                    throw new CompletionException(e);
                }
            }).thenAccept(v -> {
                logger.info("流式输出完成 - 线程ID: {}", threadId);
                
                // 发送完成事件
                try {
                    Map<String, Object> completeData = Map.of(
                        "type", "complete",
                        "threadId", threadId,
                        "message", "流式响应完成"
                    );
                    String content = objectMapper.writeValueAsString(completeData);
                    
                    ServerSentEvent<String> completeEvent = ServerSentEvent.<String>builder()
                        .id(String.valueOf(System.currentTimeMillis()))
                        .event("complete")
                        .data(content)
                        .build();
                        
                    sink.tryEmitNext(completeEvent);
                } catch (Exception e) {
                    logger.error("发送完成事件失败: {}", e.getMessage());
                }
                
                sink.tryEmitComplete();
            }).exceptionally(e -> {
                logger.error("流式输出异常 - 线程ID: {}, 错误: {}", threadId, e.getMessage(), e);
                
                // 发送错误事件
                try {
                    Map<String, Object> errorData = Map.of(
                        "type", "error",
                        "threadId", threadId,
                        "error", e.getMessage() != null ? e.getMessage() : "未知错误"
                    );
                    String content = objectMapper.writeValueAsString(errorData);
                    
                    ServerSentEvent<String> errorEvent = ServerSentEvent.<String>builder()
                        .id(String.valueOf(System.currentTimeMillis()))
                        .event("error")
                        .data(content)
                        .build();
                        
                    sink.tryEmitNext(errorEvent);
                } catch (Exception ex) {
                    logger.error("发送错误事件失败: {}", ex.getMessage());
                }
                
                sink.tryEmitError(e);
                return null;
            });
        });
    }
    
    public void shutdown() {
        executor.shutdown();
    }
}